from data_loader import *
from network import *
import time
import torch
import matplotlib.pyplot as plt

start = time.time()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

training_data, validation_data, test_data = load_data_wrapper()


def train_net(net,train_loader,optimizer,cross_entropy):
    net.train(mode=True) # 模型设置为训练模型
    epoch_loss=0
    for inputs,labels in train_loader:
        out = net(inputs)
        # 计算loss,交叉熵代价函数，out(batch,C),labels(batch)
        loss =cross_entropy(out, labels)
        epoch_loss+=loss.item()
        # 梯度清零
        optimizer.zero_grad()
        # 计算梯度
        loss.backward()
        # 修改权值
        optimizer.step()
    return epoch_loss

def validation_net(net,var_loader):
    net.eval()  # 模型设置为评估模型
    val_correct=0
    with torch.no_grad():   # 测试时不计算梯度，节省内存
        for x, y in var_loader:
            outputs = net(x)  # 输入数据包x，pre_y为其预测输出
            preds = torch.argmax(outputs, -1)
            val_correct += torch.sum(preds == y.data)
    val_acc = val_correct / len(var_loader.dataset)
    return val_acc

def test_net(net,test_loader):
    net.eval()  # 设置模型为评估模式
    test_correct = 0
    with torch.no_grad():  # 测试时不计算梯度，节省内存
        for x, y in test_loader:
            outputs = net(x)  # 输入数据包x，pre_y为其预测输出
            _, preds = torch.max(outputs, 1)
            test_correct += torch.sum(preds == y.data)
    test_acc = test_correct / len(test_loader.dataset)
    return test_acc

def train_and_test_net(net, training_data,validation_data,test_data,epochs,lr,weight_decay):
    cross_entropy = nn.CrossEntropyLoss()
    optimizer =torch.optim.Adam(net.parameters(),lr=lr, weight_decay=weight_decay)  # 优化器，添加L2正则化项

    for epoch in range(epochs):
        epoch_loss=train_net(net,training_data,optimizer,cross_entropy )  #进行模型训练
        val_accuracy=validation_net(net,validation_data)  #进行验证模型

        print(f"训练第{epoch+1}次时，训练损失值：{epoch_loss}，验证集精度{val_accuracy * 100}% ")


    # 进行测试
    test_accuracy=test_net(net, test_data)
    print(f"测试集预测准确率: {test_accuracy*100}%")


    # torch.save(net, 'net.pth')  # 保存训练好的模型
    # return val_accuracy






# 当隐藏层k=3时
# accuracy_list=[]
# for m in range(10,19):
net = Net([784,30, 10])
train_and_test_net(net,training_data,validation_data,test_data,15,0.01,0.0001)
    # accuracy_list.append(accuracy)


# x = range(10,19)
# max_acc=max(accuracy_list)
# max_index = accuracy_list.index(max_acc)
# print(f"当神经元m={10+max_index} 时，验证精度最大值为：{max_acc}")
# plt.plot(x, accuracy_list)
# plt.xlabel('m')
# plt.ylabel('val_accuracy')
# plt.show()


end = time.time()
print("运行代码所用时间： ", end - start)



